# gomoku_ai/utils/play.py
import numpy as np
from mcts.mcts import MCTS


def play_vs_human(game, net, mcts_simulations=200):
    """与AI对弈"""
    mcts = MCTS(game, net, mcts_simulations, net.device)
    board = game.get_init_board()
    player = 1  # 人类先手（黑棋）
    move_count = 0

    while True:
        game.render(board, current_player=player,
                    info_text=f"第{move_count}手 - {'您的回合(黑棋)' if player == 1 else 'AI思考中(白棋)'}")

        if player == 1:  # 人类玩家
            action = game.get_human_move(board)
        else:  # AI玩家
            canonical_board = game.get_canonical_form(board, player)
            action_probs = mcts.get_action_prob(canonical_board, temp=0)
            action = np.argmax(action_probs)

            if game.visual_board:
                row, col = divmod(action, game.board_size)
                game.visual_board.set_last_move((row, col))
                print(f"AI下在了 ({row}, {col})")

        board = game.make_move(board, player, action)
        move_count += 1

        game_result = game.get_game_ended(board)
        if game_result != 0:
            if game_result == 1:
                result_text = "恭喜！您赢了！"
            elif game_result == -1:
                result_text = "AI赢了！"
            else:
                result_text = "平局！"

            game.render(board, game_result=game_result, info_text=result_text)
            game.wait_for_continue(board, game_result)
            break

        player *= -1
